#include <os/list.h>
#include <os/mm.h>
#include <os/sched.h>
#include <pgtable.h>
#include <os/string.h>
#include <assert.h>

ptr_t memCurr = FREEMEM;

struct run {
  	struct run *next;
};

struct {
    struct run *freelist;
    uint64_t number;
}kmem;

void init_kmem(){
    kmem.freelist = NULL;
    kmem.number = 0;
}
ptr_t kmalloc_Page()
{
    // TODO:
    struct run *r=NULL;
    r=kmem.freelist;
    kmem.freelist = (kmem.freelist)->next;
    kmem.number--;
    memset((void *)r, 0, PGSIZE);
  	return (ptr_t)r;
}


void freerange(ptr_t pa_start,ptr_t pa_end){
	ptr_t p;
  	p = (ptr_t)PGROUNDUP((uint64_t)pa_start);
  	for(; p + PGSIZE <= (ptr_t)pa_end; p += PGSIZE)
    	freePage((ptr_t)p);
}

void freePage(ptr_t baseAddr)
{
  	memset((void *)baseAddr, 0, PGSIZE);
  	struct run *r;
  	r = (struct run *)baseAddr;
	r->next = kmem.freelist;
    kmem.freelist = r;
    kmem.number++;
}

void map_kpage(uint64_t vpn, uint64_t pa, PTE *pgdir,int level)
{
    // TODO: map va to pa
    pgdir[vpn]=(uint64_t)PA2PTE(pa);
    if(level == 2){
        pgdir[vpn]+=KPGTABLE_FLAG;
    }else
        pgdir[vpn]+=KPGDIR_FLAG;
}

void *kmalloc(size_t size)
{
    // TODO(if you need):
}


/* this is used for mapping kernel virtual address into user page table */
void share_pgtable(uintptr_t dest_pgdir, uintptr_t src_pgdir)
{
    // TODO:
}

/* allocate physical page for `va`, mapping it into `pgdir`,
   return the kernel virtual address for the page.
   */
uintptr_t alloc_page_helper(uintptr_t va, uintptr_t pgdir)
{
    // TODO:
}

uint64_t init_pgtable(uint64_t *pgtable,uint64_t* kernel_sp ){
    PTE *kpgdir_1st = kmalloc_Page();
    PTE *kpgdir_2nd = kmalloc_Page();
    map_kpage(0x101,kva2pa(kpgdir_2nd),kpgdir_1st,1);//for 2nd kernel page table direction
    uint64_t p;
    for(p=0;p+LARGE_PAGE_SIZE <= 0x20000000 ; p+=LARGE_PAGE_SIZE)
        map_kpage((0x10000000+p)>>LARGE_PAGE_SHIFT,0x50000000+p,kpgdir_2nd,2);
    *pgtable = (uint64_t)kpgdir_1st;
    *kernel_sp = kmalloc_Page();
    return kpgdir_1st;
    //initial 1st pgdir
    //map kernel address space
}

// Look up a virtual address, return the physical address,
// or 0 if not mapped.
// Can only be used to look up user pages.
// pgtable is kva
uint64_t walkaddr(uint64_t va, PTE* pgtable){
	uint32_t vpn2 = GETVPN2(va);
	uint32_t vpn1 = GETVPN1(va);
	uint32_t vpn0 = GETVPN0(va);
	PTE * pgtable_3rd;
    PTE * pgtable_2nd;
    PTE * page ;
	if(pgtable[vpn2]%2 == 0){
		pgtable_2nd = kmalloc_Page();
		pgtable_3rd = kmalloc_Page();
        page = kmalloc_Page();
		uint64_t pa = kva2pa(page);
		pgtable[vpn2] = (uint64_t)PA2PTE(kva2pa(pgtable_2nd))+UPGDIR_FLAG;
		pgtable_2nd[vpn1] = (uint64_t)PA2PTE(kva2pa(pgtable_3rd))+UPGDIR_FLAG;
		pgtable_3rd[vpn0] = (uint64_t)PA2PTE(pa) + UPGTABLE_FLAG;
        return page;
	}else {
        pgtable_2nd = (PTE *)pa2kva(PTE2PA(pgtable[vpn2]));
        if(pgtable_2nd[vpn1]%2 == 0){
            pgtable_3rd = kmalloc_Page();
            page = kmalloc_Page();
            uint64_t pa = kva2pa(page);
            pgtable_2nd[vpn1] = (uint64_t)PA2PTE(kva2pa(pgtable_3rd))+UPGDIR_FLAG;
            pgtable_3rd[vpn0] = (uint64_t)PA2PTE(pa) + UPGTABLE_FLAG;
        }else {
            pgtable_3rd = (PTE *)pa2kva(PTE2PA(pgtable_2nd[vpn1]));
            if(pgtable_3rd[vpn0]%2 == 0){
                page = kmalloc_Page();
                uint64_t pa = kva2pa(page);
                pgtable_3rd[vpn0] = (uint64_t)PA2PTE(pa) + UPGTABLE_FLAG;
            }
        }
    }
    page = (PTE *)pa2kva(PTE2PA(pgtable_3rd[vpn0]));
    return page;
}

uint64_t findupage(uint64_t pa){
    PTE * pgtable = current_running[current_cpu]->pgtable;
    int vpn2=0,vpn1=0,vpn0=1;
    while(pgtable[vpn2]%2 == 0 )
        vpn2++;
    PTE * pgtable_2nd = (PTE *)pa2kva(PTE2PA(pgtable[vpn2]));
    while(pgtable_2nd[vpn1]%2 == 0)
        vpn1++;
    PTE * pgtable_3rd = (PTE *)pa2kva(PTE2PA(pgtable_2nd[vpn1]));
    while(pgtable_3rd[vpn0]%2 == 1)
        vpn0++;
    pgtable_3rd[vpn0] = (uint64_t)PA2PTE(pa) + UPGTABLE_FLAG;
    uint64_t va;
    va = ((uint64_t)vpn2<<30)|((uint64_t)vpn1<<21)|((uint64_t)vpn0<<12);
    return va;
}
void freeupage(uint64_t va){
    uint32_t vpn2 = GETVPN2(va);
    uint32_t vpn1 = GETVPN1(va);
    uint32_t vpn0 = GETVPN0(va);
    PTE * pgtable = current_running[current_cpu]->pgtable;
    if(pgtable[vpn2]%2 == 0)
        return;
    PTE * pgtable_2nd = (PTE *)pa2kva(PTE2PA(pgtable[vpn2]));
    if(pgtable_2nd[vpn1]%2 == 0)
        return;
    PTE * pgtable_3rd = (PTE *)pa2kva(PTE2PA(pgtable_2nd[vpn1]));
    pgtable_3rd[vpn0] = 0;
    return;
}

